package com.zx.gateway.controller

import com.zx.boot.enums.ModuleType
import com.zx.boot.util.Slf4j.Companion.LOGGER
import com.zx.cloud.nacos.control.ZxCloudNacosControlCenter
import org.apache.commons.lang3.StringUtils
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.core.io.buffer.DataBuffer
import org.springframework.http.HttpHeaders
import org.springframework.http.HttpMethod
import org.springframework.web.bind.annotation.RequestMapping
import org.springframework.web.bind.annotation.RestController
import org.springframework.web.reactive.function.BodyInserters
import org.springframework.web.reactive.function.client.ClientResponse
import org.springframework.web.reactive.function.client.WebClient
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Mono
import java.time.Duration

@RestController
class ApiProxyController {
    companion object {
        private val webClient = WebClient.create()
        private const val BACKEND_SERVICE_TIMEOUT_IN_MILLIS = 60 * 60 * 1000
    }

    @Autowired
    private lateinit var cloudNacosControlCenter: ZxCloudNacosControlCenter


    @RequestMapping("/**")
    fun proxyRequest(exchange: ServerWebExchange): Mono<Void> {
        val uri = exchange.request.uri.toString().replaceFirst("/gateway", "")
        if (LOGGER.isDebugEnabled){
            LOGGER.debug("before ==> {}", uri)
        }
        val httpMethod = exchange.request.method!!
        val index = StringUtils.ordinalIndexOf(uri, "/", 3)
        val index1 = StringUtils.ordinalIndexOf(uri, "/", 4)
        val index2 = StringUtils.ordinalIndexOf(uri, "/", 5)
        val service = uri.substring(index + 1, index1)
        val module = uri.substring(index1 + 1, index2)
        val servers = cloudNacosControlCenter.services()
        if (servers.none { it.id == service }) {
            throw RuntimeException("请检查 ${service} 健康情况！")
        }
        if (servers.filter { it.id == service }[0].getModule(module).type == ModuleType.INNER) {
            throw RuntimeException("请求路径异常！")
        }
        val host = cloudNacosControlCenter.serviceUrl(service)
        val servletPath = uri.substring(index1)
        val forward = host + servletPath
        if (LOGGER.isDebugEnabled){
            LOGGER.debug("after ==> {}", forward)
        }
        val reqBodySpec = webClient.method(httpMethod).uri(forward).headers { httpHeaders: HttpHeaders ->
            httpHeaders.addAll(exchange.request.headers)
            httpHeaders.remove("HOST")
        }
        val reqHeadersSpec = if (requireHttpBody(httpMethod)) {
            reqBodySpec.body(BodyInserters.fromDataBuffers(exchange.request.body))
        } else {
            reqBodySpec
        }

        // 调用后端服务
        return reqHeadersSpec.exchange().timeout(Duration.ofMillis(BACKEND_SERVICE_TIMEOUT_IN_MILLIS.toLong()))
                .flatMap { backendResponse: ClientResponse ->
                    exchange.response.statusCode = backendResponse.statusCode()
                    exchange.response.headers.putAll(backendResponse.headers().asHttpHeaders())
                    exchange.response.writeWith(backendResponse.bodyToFlux(DataBuffer::class.java))
                }
    }

    private fun requireHttpBody(method: HttpMethod): Boolean {
        return HttpMethod.POST == method || HttpMethod.PUT == method || HttpMethod.PATCH == method || HttpMethod.TRACE == method
    }
}
